# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/13 15:59 
@Author  : ZhangShenao 
@File    : 2.使用SequentialChain顺序链.py 
@Desc    : 使用SequentialChain顺序链
"""
import os

import dotenv
from langchain.chains.llm import LLMChain
from langchain.chains.sequential import SequentialChain
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

# 加载环境变量
dotenv.load_dotenv()

# 创建LLM
llm = ChatOpenAI(model_name='gpt-3.5-turbo',
                 openai_api_base=os.getenv('OPENAI_API_BASE'),
                 temperature=0.8
                 )

# 构造第1个Chain:生成专业的鲜花介绍
introduction_template = '''
你是一位资深的植物学家。给定花的名称和类型，请你为这种花写一个200字左右的介绍。
花名: {name}
颜色: {color}
植物学家: 这是关于上述花的介绍：
'''
introduction_prompt = ChatPromptTemplate.from_template(introduction_template)
introduction_chain = LLMChain(llm=llm, prompt=introduction_prompt, output_key='introduction')

# 构造第2个Chain:根据鲜花的介绍,撰写鲜花的评论
review_template = '''
你是一位专业的鲜花评论家。给定一种花的介绍，请你为这种花写一篇200字左右的评论。
鲜花介绍:
{introduction}
花评人对上述花的评论:
'''
review_prompt = ChatPromptTemplate.from_template(review_template)
review_chain = LLMChain(llm=llm, prompt=review_prompt, output_key='review')

# 构造第3个Chain:根据鲜花的介绍和评论,发表一篇自媒体文章
post_template = '''
你是一家花店的社交媒体经理。给定一种花的介绍和评论，请你为这种花写一篇社交媒体的帖子，字数为300字左右。
鲜花介绍:
{introduction}
花评人对上述花的评论:
{review}
社交媒体帖子:
'''
post_prompt = ChatPromptTemplate.from_template(post_template)
post_chain = LLMChain(llm=llm, prompt=post_prompt, output_key='post')

# 构造SequentialChain,顺序执行上面的3个Chain
chain = SequentialChain(
    chains=[introduction_chain, review_chain, post_chain],
    input_variables=['name', 'color'],
    output_variables=['introduction', 'review', 'post'],
    verbose=True
)

# 传递输入参数,执行Chain,获取执行结果
output = chain.invoke({
    'name': '玫瑰花',
    'color': '白色',
})
print(output)
